package com.lun.medium;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import com.lun.util.BinaryTree.TreeNode;

public class AllNodesDistanceKInBinaryTree {
	
	//方法一：我写的
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
    	LinkedList<TreeNode> path = findTarget(root, target);
    	List<Integer> result = new ArrayList<>();
    	findUp(k, path, result);
    	findDown(k, path.peekLast(), result);
    	return result;
    }

	public void findUp(int k, LinkedList<TreeNode> path, List<Integer> result) {
		int pathTargetIndex = path.size() - 1;
    	int pathInitCount = 1;
    	int currentIndex = pathTargetIndex - 1;
    	int remainder = k - pathInitCount;//剩余步数
    	for(;currentIndex >= 0 && remainder >= 0; // 
    			pathInitCount++, currentIndex--, remainder = k - pathInitCount) {

    		TreeNode currentNode = path.get(currentIndex);
    		if(remainder > 0) {
    			TreeNode nextNode = path.get(currentIndex + 1);
    			
    			//在右子树查找
    			LinkedList<TreeNode> queue = new LinkedList<>();
    			queue.offer(currentNode);
    			int count = 0;
    	    	while(!queue.isEmpty()) {
    	    		
    	    		for(int size = queue.size(); size > 0; size--) {
    	    			TreeNode node = queue.poll();
    	    			
    	    			if(currentNode.left == nextNode) {
    	    				//只查右子树
    	    				
        	    			if(node.left != null && node.left != nextNode)
        	    				queue.offer(node.left);
        	    			
        	    			if(node.right != null)
        	    				queue.offer(node.right);
    	    				
    	    			}else {
    	    				
        	    			if(node.left != null)//刚开始一层有点特殊
        	    				queue.offer(node.left);
        	    			
        	    			if(node.right != null && node.right != nextNode)
        	    				queue.offer(node.right);
    	    			}
    	    			
    	    		}
    	    		
    	    		if(++count == remainder) {
    	    			queue.stream().forEach(a->result.add(a.val));
    	    			break;
    	    		}
    	    	}
    		}else if(remainder == 0) {
    			result.add(currentNode.val);
    		}
    	}
	}
	
	public void findDown(int k, TreeNode startNode, List<Integer> result) {
		
		if(k == 0) {//findUp不用，只一个即可
			result.add(startNode.val);
			return;
		}
			
		LinkedList<TreeNode> queue = new LinkedList<>();
    	queue.offer(startNode);
    	int count = 0;
    	while(!queue.isEmpty()) {
    		
    		for(int size = queue.size(); size > 0; size--) {
    			TreeNode node = queue.poll();
    			
    			if(node.left != null)
    				queue.offer(node.left);
    			
    			if(node.right != null)
    				queue.offer(node.right);
    		}
    		
    		if(++count == k) {
    			queue.stream().forEach(a->result.add(a.val));
    			break;
    		}
    	}
	}
	
    
    //目标节点到跟节点的路径
    @SuppressWarnings("unchecked")
	public LinkedList<TreeNode> findTarget(TreeNode root, TreeNode target){
    	
    	LinkedList<Object[]> stack = new LinkedList<>();
    	
    	stack.push(new Object[] {root, new LinkedList<TreeNode>()});
    	
    	while(!stack.isEmpty()) {
    		Object[] arr = stack.pop();
    		TreeNode node = (TreeNode)arr[0];
    		LinkedList<TreeNode> path = (LinkedList<TreeNode>)arr[1];
    	
    		path.add(node);
    		
    		if(node.val == target.val)
    			return path;
    		
    		if(node.left != null && node.right == null) {
    			stack.push(new Object[] {node.left, path});
    		}
    		
    		if(node.left == null && node.right != null) {
    			stack.push(new Object[] {node.right, path});
    		}
    		
    		if(node.left != null && node.right != null) {
    			stack.push(new Object[] {node.right, new LinkedList<TreeNode>(path)});
    			stack.push(new Object[] {node.left, path});
    		}
    	}
    	
    	return new LinkedList<TreeNode>();
    }
    
    //方法二：别人写的
    public static class Solution1{
        Map<TreeNode, Integer> map = new HashMap<>();
        
        public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
            List<Integer> res = new LinkedList<>();
            find(root, target);
            dfs(root, K, map.get(root), res);
            return res;
        }
        
        // find target node first and store the distance in that path that we could use it later directly
        private int find(TreeNode root, TreeNode target) {
            if (root == null) return -1;
            if (root.val == target.val) {
                map.put(root, 0);
                return 0;
            }
            int left = find(root.left, target);
            if (left >= 0) {
                map.put(root, left + 1);
                return left + 1;
            }
    		int right = find(root.right, target);
    		if (right >= 0) {
                map.put(root, right + 1);
                return right + 1;
            }
            return -1;
        }
        
        private void dfs(TreeNode root, int K, int length, List<Integer> res) {
            if (root == null) return;
            if (map.containsKey(root)) length = map.get(root);
            if (length == K) res.add(root.val);
            dfs(root.left, K, length + 1, res);
            dfs(root.right, K, length + 1, res);
        }
    }
    
    //方法三：别人写的2，方法二缝合版
    public static class Solution2{
    	
        public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
            List<Integer> res = new LinkedList<>();
            if (K == 0) {
                res.add(target.val);
            } else {
                dfs(res, root, target.val, K ,0);
            }
            return res;
        }
        
        private int dfs(List<Integer> res, TreeNode node, int target, int K, int depth) {
            if (node == null) return 0;
            
            if (depth == K) {
                res.add(node.val);
                return 0;
            }
            
            int left, right;
            if (node.val == target || depth > 0) {
                left = dfs(res, node.left, target, K, depth + 1);
                right = dfs(res, node.right, target, K, depth + 1);
            } else {
                left = dfs(res, node.left, target, K, depth);
                right = dfs(res, node.right, target, K, depth);
            }
            
            if (node.val == target) return 1;
            
            if (left == K || right == K) {
                res.add(node.val);
                return 0;
            }
            
            if (left > 0) {
                dfs(res, node.right, target, K, left + 1);
                return left + 1;
            }
            
            if (right > 0) {
                dfs(res, node.left, target, K, right + 1);
                return right + 1;
            }
            
            return 0;
        }
    }
    
    
}
